from multiprocessing import cpu_count
from kdtree import build_kdtree, get_closest_object_from_kdtree, kdtree_closest_point, get_all_line_values_from_kdtree
from rtree import RtreeNode, Extent, search_rtree_extent, insert
from trees import Object, get_objects_list_for_kdtree, get_objects_list_for_rtree, get_distance_between_objects, get_object_MBR

import json
import unittest
import cProfile_performance_test

class TestTrees(unittest.TestCase):
    # Проверка работы KD-дерева на объектах, состоящих только из типа Point
    def test_kdtree_for_points(self):
        with open("test_data/test_kdtree_for_points.geojson", 'r') as file:
            geojson_data = json.load(file)
        
        objects_list = get_objects_list_for_kdtree(geojson_data)
        defined_object = [
            [
                108.6328125, 
                72.86793049861396
            ]
        ]
        root = build_kdtree(objects_list)

        closest_object = []

        for value in defined_object:
            closest_object.append(kdtree_closest_point(root, value))
        
        finded_node = get_closest_object_from_kdtree(closest_object)
        finded_node['geometry']['coordinates'] = get_all_line_values_from_kdtree(root, finded_node, [])

        test_data = {
            'id' : 0, 
            'geometry' : {
                'type' : 'Point',
                'coordinates' : [
                    [
                        111.09374999999999,
                        70.25945200030638
                    ]
                ]
            },
            'distance' : 290867
        }

        self.assertEqual(finded_node, test_data)

    # Проверка работы KD-дерева на объектах, состоящих только из типа Polyline
    def test_kdtree_for_polylines(self):
        with open("test_data/test_kdtree_for_polylines.geojson", 'r') as file:
            geojson_data = json.load(file)
        
        objects_list = get_objects_list_for_kdtree(geojson_data)
        defined_object = [
            [
                55.5029296875,
                57.844750992891
            ],
            [
                55.5908203125,
                57.27904276497778
            ],
            [
                59.58984374999999,
                56.24334992410525
            ],
            [
                66.8408203125,
                60.4788788301667
            ]
        ]
        root = build_kdtree(objects_list)

        closest_object = []

        for value in defined_object:
            closest_object.append(kdtree_closest_point(root, value))
        
        finded_node = get_closest_object_from_kdtree(closest_object)
        finded_node['geometry']['coordinates'] = get_all_line_values_from_kdtree(root, finded_node, [])

        test_data = {
            'id' : 1, 
            'geometry' : {
                'type' : 'LineString',
                'coordinates' : [
                    [
                        64.775390625, 
                        52.5897007687178
                    ], 
                    [
                        58.97460937499999, 
                        52.29504228453735
                    ], 
                    [
                        60.68847656249999, 
                        54.08517342088679
                    ]
                ]
            },
            'distance' : 170912
        }

        self.assertEqual(finded_node, test_data)

    # Проверка работы KD-дерева на объектах, состоящих из типов Point и Polyline при заданном объекте типа Point
    def test_kdtree_for_points_and_polylines_with_defined_point(self):
        with open("test_data/test_kdtree_for_points_and_polylines.geojson", 'r') as file:
            geojson_data = json.load(file)
        
        objects_list = get_objects_list_for_kdtree(geojson_data)
        defined_object = [
            [
                61.5673828125,
                60.823494332539646
            ]
        ]
        root = build_kdtree(objects_list)

        closest_object = []

        for value in defined_object:
            closest_object.append(kdtree_closest_point(root, value))
        
        finded_node = get_closest_object_from_kdtree(closest_object)
        finded_node['geometry']['coordinates'] = get_all_line_values_from_kdtree(root, finded_node, [])

        test_data = {
            'id' : 5, 
            'geometry' : {
                'type' : 'LineString',
                'coordinates' : [
                    [
                        55.8984375, 
                        62.97519757003264
                    ], 
                    [
                        56.5576171875,
                        63.509375401175134
                    ], 
                    [
                        57.48046875, 
                        63.99523519297698
                    ], 
                    [
                        57.568359375, 
                        62.61356210229029
                    ], 
                    [
                        56.77734375,
                        62.512317938386914
                    ], 
                    [
                        58.00781249999999, 
                        64.84893726357947
                    ], 
                    [
                        56.99707031249999, 
                        65.2198939361321
                    ]
                ]
            },
            'distance' : 456041
        }

        self.assertEqual(finded_node, test_data)

    # Проверка работы KD-дерева на объектах, состоящих из типов Point и Polyline при заданном объекте типа Polilyne
    def test_kdtree_for_points_and_polylines_with_defined_polyline(self):
        with open("test_data/test_kdtree_for_points_and_polylines.geojson", 'r') as file:
            geojson_data = json.load(file)
        
        objects_list = get_objects_list_for_kdtree(geojson_data)
        defined_object = [
            [
                61.69921875,
                47.66538735632654
            ],
            [
                57.1728515625,
                50.261253827584724
            ],
            [
                58.71093750000001,
                51.31688050404585
            ],
            [
                55.107421875,
                51.754240074033525
            ],
            [
                55.8984375,
                54.7246201949245
            ]
        ]
        root = build_kdtree(objects_list)

        closest_object = []

        for value in defined_object:
            closest_object.append(kdtree_closest_point(root, value))
        
        finded_node = get_closest_object_from_kdtree(closest_object)
        finded_node['geometry']['coordinates'] = get_all_line_values_from_kdtree(root, finded_node, [])

        test_data = {
            'id' : 1, 
            'geometry' : {
                'type' : 'LineString',
                'coordinates' : [
                    [
                        64.775390625, 
                        52.5897007687178
                    ], 
                    [
                        60.68847656249999, 
                        54.08517342088679
                    ], 
                    [
                        58.97460937499999, 
                        52.29504228453735
                    ]
                ]
            },
            'distance' : 63471
        }

        self.assertEqual(finded_node, test_data)

    # Проверка работы R-дерева на объектах, состоящих только из типа Point
    def test_rtree_for_points(self):
        with open("test_data/test_rtree_for_points.geojson", 'r') as file:
            geojson_data = json.load(file)
        objects_list = get_objects_list_for_rtree(geojson_data)
        M = 3
        root = RtreeNode(M+1)

        extents = [Extent(obj.MBR[0], obj.MBR[1], obj.MBR[2], obj.MBR[3], obj.id, obj.geometry) for obj in objects_list]

        for e in extents:
            n = search_rtree_extent(root, e)
            insert(n, e)

        coordinates = [
            [
                64.896240234375,
                59.130863097255904
            ]
        ]
        
        MBR = get_object_MBR(coordinates)
        defined_object = Object(geometry = coordinates, MBR = Extent(MBR[0], MBR[1], MBR[2], MBR[3]))

        finded_node = search_rtree_extent(root, defined_object.MBR)

        distance_list = get_distance_between_objects(finded_node.entries, defined_object)

        test_data_distance = [115612, 45889]
        test_data_geometry = [
            {
                'id' : 17,
                'type' : 'Point',
                'coordinates' : [
                        65.91796875,
                        59.58997510773555
                ]
            },
            {
                'id': 5,
                'type' : 'Point',
                'coordinates' : [
                        64.57763671875, 
                        58.516651799363785
                ]
            }
        ]

        for i, ent in enumerate(finded_node.entries):
            self.assertEqual({'id' : ent.id, 'type' : ent.geometry['type'], 'coordinates' : ent.geometry['coordinates']}, test_data_geometry[i])
            self.assertEqual(distance_list[i], test_data_distance[i])

    # Проверка работы R-дерева на объектах, состоящих только из типа LineString
    def test_rtree_for_polylines(self):
        with open("test_data/test_rtree_for_polylines.geojson", 'r') as file:
            geojson_data = json.load(file)
        objects_list = get_objects_list_for_rtree(geojson_data)
        M = 4
        root = RtreeNode(M+1)

        extents = [Extent(obj.MBR[0], obj.MBR[1], obj.MBR[2], obj.MBR[3], obj.id, obj.geometry) for obj in objects_list]

        for e in extents:
            n = search_rtree_extent(root, e)
            insert(n, e)

        coordinates = [
            [
                76.025390625,
                56.70450561416937
            ],
            [
                64.423828125,
                60.37042901631508
            ],
            [
                61.04003906249999,
                65.01650627048231
            ],
            [
                47.8564453125,
                66.8265202749748
            ]
        ]
        
        MBR = get_object_MBR(coordinates)
        defined_object = Object(geometry = coordinates, MBR = Extent(MBR[0], MBR[1], MBR[2], MBR[3]))

        finded_node = search_rtree_extent(root, defined_object.MBR)

        distance_list = get_distance_between_objects(finded_node.entries, defined_object)

        test_data_distance = [911050, 138334, 292811]
        test_data_geometry = [
            {
                'id': 5,
                'type' : 'LineString',
                'coordinates' : [
                        [59.4140625, 60.84491057364912], [59.0185546875, 60.06484046010452], [58.4033203125, 59.5343180010956], [59.150390625, 59.108308258604964], [59.2822265625, 58.6769376725869], [59.80957031249999, 58.47072082411973], [58.97460937499999, 58.19387126497797]
                ]
            },
            {
                'id' : 4,
                'type' : 'LineString',
                'coordinates' : [
                        [66.70898437499999, 58.6769376725869], [66.09375, 58.60833366077633], [66.181640625, 58.0546319113729], [65.6982421875, 57.75107598132104], [65.3466796875, 57.9148477670092], [64.7314453125, 57.562995459387146]
                ]
            },
            {
                'id' : 0,
                'type' : 'LineString',
                'coordinates' : [
                        [62.40234375, 54.059387886623576], [64.64355468749999, 54.36775852406841], [65.654296875, 54.7246201949245], [66.3134765625, 54.57206165565852], [67.8515625, 54.95238569063361]
                ]
            }
        ]

        for i, ent in enumerate(finded_node.entries):
            self.assertEqual({'id' : ent.id, 'type' : ent.geometry['type'], 'coordinates' : ent.geometry['coordinates']}, test_data_geometry[i])
            self.assertEqual(distance_list[i], test_data_distance[i])

    # Проверка работы R-дерева на объектах, состоящих из типов Point и Polyline при заданном объекте типа Point
    def test_rtree_for_points_and_polylines_with_defined_point(self):
        with open("test_data/test_rtree_for_points_and_polylines.geojson", 'r') as file:
            geojson_data = json.load(file)
        objects_list = get_objects_list_for_rtree(geojson_data)
        M = 6
        root = RtreeNode(M+1)

        extents = [Extent(obj.MBR[0], obj.MBR[1], obj.MBR[2], obj.MBR[3], obj.id, obj.geometry) for obj in objects_list]

        for e in extents:
            n = search_rtree_extent(root, e)
            insert(n, e)

        coordinates = [
            [
                65.126953125,
                60.65164736580915
            ]
        ]
        
        MBR = get_object_MBR(coordinates)
        defined_object = Object(geometry = coordinates, MBR = Extent(MBR[0], MBR[1], MBR[2], MBR[3]))

        finded_node = search_rtree_extent(root, defined_object.MBR)

        distance_list = get_distance_between_objects(finded_node.entries, defined_object)

        test_data_distance = [939773, 170240, 1385106, 269903]
        test_data_geometry = [
            {
                'id' : 23,
                'type' : 'Point',
                'coordinates' : [
                        73.3447265625, 55.00282580979323
                ]
            },
            {
                'id': 16,
                'type' : 'Point',
                'coordinates' : [
                        65.56640625, 57.136239319177434
                ]
            },
            {
                'id': 18,
                'type' : 'LineString',
                'coordinates' : [
                        [48.9990234375, 53.904338156274704], [48.2958984375, 54.49556752187406], [49.2626953125, 55.00282580979323], [49.5263671875, 55.25407706707272], [51.37207031249999, 55.50374985927514], [52.294921875, 55.825973254619015], [52.91015625, 55.87531083569679]
                ]
            },
            {
                'id' : 2,
                'type' : 'LineString',
                'coordinates' : [
                        [59.765625, 63.99523519297698], [59.2822265625, 63.194018438087845], [59.6337890625, 62.552856958572896], [59.45800781249999, 61.91827102335593], [60.029296875, 61.87687021463305], [60.16113281250001, 61.64816245852389], [60.732421875, 61.75233128411639], [62.70996093749999, 61.14323525084058]
                ]
            },
        ]

        for i, ent in enumerate(finded_node.entries):
            self.assertEqual({'id' : ent.id, 'type' : ent.geometry['type'], 'coordinates' : ent.geometry['coordinates']}, test_data_geometry[i])
            self.assertEqual(distance_list[i], test_data_distance[i])
            
    # Проверка работы R-дерева на объектах, состоящих из типов Point и Polyline при заданном объекте типа Polilyne
    def test_rtree_for_points_and_polylines_with_defined_polyline(self):
        with open("test_data/test_rtree_for_points_and_polylines.geojson", 'r') as file:
            geojson_data = json.load(file)
        objects_list = get_objects_list_for_rtree(geojson_data)
        M = 5
        root = RtreeNode(M+1)

        extents = [Extent(obj.MBR[0], obj.MBR[1], obj.MBR[2], obj.MBR[3], obj.id, obj.geometry) for obj in objects_list]

        for e in extents:
            n = search_rtree_extent(root, e)
            insert(n, e)

        coordinates = [
            [
                77.87109375,
                68.43151284537514
            ],
            [
                68.8623046875,
                65.20146842127558
            ],
            [
                61.787109375,
                64.90491004905083
            ],
            [
                66.533203125,
                59.88893689676585
            ],
            [
                55.3271484375,
                50.764259357116465
            ],
            [
                42.6708984375,
                52.669720383688166
            ]
        ]
        
        MBR = get_object_MBR(coordinates)
        defined_object = Object(geometry = coordinates, MBR = Extent(MBR[0], MBR[1], MBR[2], MBR[3]))

        finded_node = search_rtree_extent(root, defined_object.MBR)

        distance_list = get_distance_between_objects(finded_node.entries, defined_object)

        test_data_distance = [139191, 731526, 1516744]
        test_data_geometry = [
            {
                'id' : 12,
                'type' : 'Point',
                'coordinates' : [
                        76.728515625, 66.10716955858042
                ]
            },
            {
                'id': 15,
                'type' : 'Point',
                'coordinates' : [
                        76.5966796875, 60.8663124746226
                ]
            },
            {
                'id' : 14,
                'type' : 'Point',
                'coordinates' : [
                        75.41015624999999, 63.15435519659187
                ]
            }
        ]

        for i, ent in enumerate(finded_node.entries):
            self.assertEqual({'id' : ent.id, 'type' : ent.geometry['type'], 'coordinates' : ent.geometry['coordinates']}, test_data_geometry[i])
            self.assertEqual(distance_list[i], test_data_distance[i])


if __name__ == "__main__":
    test_obj = TestTrees()
    cProfile_performance_test.performance_info("test_obj.test_kdtree_for_points")
    cProfile_performance_test.performance_info("test_obj.test_kdtree_for_polylines")
    cProfile_performance_test.performance_info("test_obj.test_kdtree_for_points_and_polylines_with_defined_point")
    cProfile_performance_test.performance_info("test_obj.test_kdtree_for_points_and_polylines_with_defined_polyline")
    cProfile_performance_test.performance_info("test_obj.test_rtree_for_points")
    cProfile_performance_test.performance_info("test_obj.test_rtree_for_polylines")
    cProfile_performance_test.performance_info("test_obj.test_rtree_for_points_and_polylines_with_defined_point")
    cProfile_performance_test.performance_info("test_obj.test_rtree_for_points_and_polylines_with_defined_polyline")